#include "RUDP.h"
#include "Serialize.h"

RUDP::RUDP(unsigned int initRTO, unsigned short fast_rsend, unsigned short MTU, unsigned char conv) :
	conv(conv),
	RTO(initRTO),
	MTU(MTU),
	fast_rsend(fast_rsend),
	send_nxt(0x0),
	recv_nxt(0x0),
	max_seq(0xffffffff),
	seq_rounded(false),
	tp(std::chrono::time_point_cast<std::chrono::milliseconds>(std::chrono::system_clock::now())),
	buffer(new char[MTU * 2]),
	enable(true)
								
{
}

RUDP::~RUDP()
{
	for (auto iter = send_buf.begin(); iter != send_buf.end(); iter++)
		releaseFragment(iter->second);
	send_buf.clear();

	for (auto iter = recv_queue.begin(); iter != recv_queue.end(); iter++)
		releaseFragment(iter->second);
	recv_queue.clear();

	while (!recv_buf.empty())
	{
		releaseFragment(recv_buf.front());
		recv_buf.pop();
	}

	while (!ack_queue.empty())
	{
		releaseFragment(ack_queue.front());
		ack_queue.pop();
	}

	delete[] buffer;
}

void RUDP::SetOutput(std::function<void(const char *buf, int size)> send_func)
{
	this->send_func = send_func;
}

void RUDP::Input(char *msg, unsigned int size, const char *ip)
{
	if (!enable)
		return;

	unsigned char conv = 0;
	unsigned char msgType = 0;
	char *ptr = msg;

	while (size > ptr - msg)
	{
		ptr = Serialize::Deserialize8(ptr, conv);

		if (conv != this->conv)
			return;

		ptr = Serialize::Deserialize8(ptr, msgType);

		switch (msgType)
		{
		case PSH:
			ptr = PSHHandler(ptr, msg + size);
			if (ptr == nullptr)
				return;
			break;
		case ACK:
			ptr = ACKHandler(ptr);
			break;
		default:
			break;
		}
	}
}

bool RUDP::Send(const char *msg, unsigned short size)
{
	if (!enable)
		return false;

	//	serialize fragment header
	if (size + HeaderLen > MTU)
		return false;
	char *buf = new char[size + HeaderLen - 4];
	char *ptr = buf;
	//	communication id
	ptr = Serialize::Serialize8(ptr, conv);
	//	message type
	ptr = Serialize::Serialize8(ptr, PSH);
	//	sent sequence
	ptr = Serialize::Serialize32(ptr, send_nxt);
	//	return time stamp
	ptr = Serialize::Serialize32(ptr, this->Clock());
	//	playload length
	ptr = Serialize::Serialize16(ptr, size);
	memcpy(ptr, msg, size);
	ptr += size;

	Fragment *newMSG = new Fragment(buf, ptr - buf, Clock() + RTO);
	send_buf[send_nxt] = newMSG;
	if (newMSG->size > 0)
		send_func(newMSG->buf, newMSG->size);

	send_nxt++;
	return true;
}

bool RUDP::Recv(std::queue<UserData *> &msg)
{
	if (recv_buf.empty())
		return false;

	while (!recv_buf.empty())
	{
		msg.push(recv_buf.front());
		recv_buf.pop();
	}

	return true;
}

unsigned int RUDP::PendingSendCount()
{
	return send_buf.size();
}

unsigned int RUDP::PendingRecvCount()
{
	return recv_buf.size();
}

void RUDP::Disable()
{
	enable = false;
}

bool RUDP::isEnable()
{
	return enable;
}

void RUDP::Update()
{
	if (!enable)
		return;

	char *ptr = buffer;
	unsigned int predict_size = 0;
	//	pack ack fragment
	while (!ack_queue.empty())
	{
		unsigned int msg_size = ack_queue.front()->size;
		predict_size = ptr - buffer + msg_size;
		if (MTU >= predict_size)
		{
			memcpy(ptr, ack_queue.front()->buf, msg_size);
			ptr += msg_size;
			releaseFragment(ack_queue.front());
			ack_queue.pop();
		}
		else
			break;
	}

	//	pack data fragment by timer
	for (auto iter = send_buf.begin(); iter != send_buf.end(); iter++)
	{
		int ts = this->Clock() - iter->second->TS;
		if (ts >= 0)
		{
			predict_size = ptr - buffer + iter->second->size;
			if (MTU >= predict_size)
			{
				UpdateFragmentTS(iter->second->buf);
				memcpy(ptr, iter->second->buf, iter->second->size);
				ptr += iter->second->size;
			}
			iter->second->TS = this->Clock() + RTO;
			iter->second->tri_fastrsnd = false;
		}
	}

	//	send all of data to endpoint
	int size = ptr - buffer;
	if (size > 0)
	{
		send_func(buffer, size);
	}
}

void RUDP::releaseFragment(Fragment *msg)
{
	//	release memory
	delete[] msg->buf;
	delete msg;
}

char *RUDP::PSHHandler(char *msg, char *end_msg)
{
	unsigned int seq = 0;
	unsigned int ts = 0;
	unsigned short playloadLen = 0;
	char *ptr = msg;
	//	deserialize fragment header
	ptr = Serialize::Deserialize32(ptr, seq);
	ptr = Serialize::Deserialize32(ptr, ts);
	ptr = Serialize::Deserialize16(ptr, playloadLen);

	if (playloadLen > end_msg - ptr)
		return nullptr;

	//	expect sequence 0 to 0xffffffff >> 3
	UpdateRecvRound();
	if (InRange(seq, 0, 0xffffffff >> 3) && !seq_rounded)
	{
		if (recv_queue.find(seq) == recv_queue.end())
			recv_queue[seq] = NewFragment(ptr, playloadLen);
	}
	// receive new Fragment
	if (InRange(seq, 0, max_seq))
	{
		PushFragment(seq, ptr, playloadLen);
	}
	ParseACKFragment(seq, ts);
	return ptr + playloadLen;
}

char *RUDP::ACKHandler(char *msg)
{
	unsigned int seq;
	unsigned int unack;
	unsigned int ts;
	unsigned int curr;
	long long RTT = 0;
	unsigned int predict_size;
	char *ptr = msg;
	ptr = Serialize::Deserialize32(ptr, seq);
	ptr = Serialize::Deserialize32(ptr, unack);
	ptr = Serialize::Deserialize32(ptr, ts);

	//	delete previous fragment
	long long ack = 0;
	ack = (long long)unack - 1;
	bool suces;
	DelSendBuf(seq, suces);
	unsigned int size = send_buf.size();
	for (unsigned int i = 0; size > i; i++)
	{
		if (ack == 0)
			ack = 0xffffffff;
		if (0 > ack)
			break;
		DelSendBuf(ack, suces);
		ack--;
	}

	//	recalculate RTO by RTT
	curr = this->Clock();
	RTT = (long long)curr - ts;
	if (0 > RTT)
		RTT = (0xffffffff - ts) + curr;
	RTO = RTT + RTO / 2;

	// fast resend, when skip over count of fast_rsend value
	long long need_rsend = 0;
	need_rsend = (long long)seq - unack;
	if (need_rsend < 0)
	{
		if (InRange(seq, 0, 0xffffffff >> 3) && InRange(unack, (0xffffffff - (0xffffffff >> 3)), 0xffffffff))
			need_rsend = (0xffffffff - unack) + seq;
	}
	if (need_rsend > fast_rsend)
	{
		if (!InRange(need_rsend, fast_rsend, 0xffffffff >> 3))
			return ptr;

		char *ptr = buffer;
		for (unsigned int ack = unack; ack != seq; ack++)
		{
			auto iter = send_buf.find(ack);
			if (iter != send_buf.end())
			{
				if (!iter->second->tri_fastrsnd)
				{
					predict_size = ptr - buffer + iter->second->size;
					if (MTU >= predict_size)
					{
						UpdateFragmentTS(iter->second->buf);
						memcpy(ptr, iter->second->buf, iter->second->size);
						ptr += iter->second->size;
						iter->second->TS = this->Clock() + RTO;
						iter->second->tri_fastrsnd = true;
					}
					else
						break;
				}
			}
		}

		//	send all of data to endpoint
		int size = ptr - buffer;
		if (size > 0)
		{
			send_func(buffer, size);
		}
	}

	return ptr;
}

void RUDP::PushFragment(unsigned int seq, char *data, unsigned int size)
{
	//	push new fragment to recv_buf
	CheckRecvQueue();
	if (recv_nxt == seq)
	{
		recv_buf.push(NewFragment(data, size));
		recv_nxt++;
		UpdateRecvRound();
	}
	else if (seq > recv_nxt)
	{
		if (recv_queue.find(seq) == recv_queue.end())
			recv_queue[seq] = NewFragment(data, size);
	}
	CheckRecvQueue();
}

void RUDP::ParseACKFragment(unsigned int seq, unsigned int ts)
{
	//	serialize fragment header
	char *buf = new char[HeaderLen - 2];
	char *ptr = buf;
	//	communication id
	ptr = Serialize::Serialize8(ptr, conv);
	//	message type
	ptr = Serialize::Serialize8(ptr, ACK);
	//	received sequence
	ptr = Serialize::Serialize32(ptr, seq);
	//	received next sequence
	ptr = Serialize::Serialize32(ptr, recv_nxt);
	//	return time stamp
	ptr = Serialize::Serialize32(ptr, ts);
	ack_queue.push(new Fragment(buf, ptr - buf));
}

void RUDP::UpdateFragmentTS(char *msg)
{
	//	update time stamp for RTT
	char *ptr = msg;
	const unsigned short TSIndex = 6;
	Serialize::Serialize32(&ptr[TSIndex], this->Clock());
}

void RUDP::CheckRecvQueue()
{
	//	arrange received fragments
	while (!recv_queue.empty())
	{
		auto msg = recv_queue.find(recv_nxt);
		if (msg != recv_queue.end())
		{
			recv_buf.push(msg->second);
			recv_queue.erase(msg);
			recv_nxt++;
			UpdateRecvRound();
		}
		else
			break;
	}
}

void RUDP::UpdateRecvRound()
{
	//	when recv_nxt overflowed, set rec_round as true
	if (recv_nxt == 0)
	{
		seq_rounded = true;
		max_seq = 0xffffffff >> 1;
	}
	//	assume all of previous fragment are cleared and confirmed
	else if (recv_nxt >= 0xffffffff >> 2)
	{
		seq_rounded = false;
		max_seq = 0xffffffff;
	}
}

std::map<unsigned int, RUDP::Fragment *>::iterator RUDP::DelSendBuf(unsigned int ack, bool &succes)
{
	auto sendBufIter = send_buf.find(ack);
	if (sendBufIter != send_buf.end())
	{
		succes = true;
		releaseFragment(sendBufIter->second);
		return send_buf.erase(sendBufIter);
	}
	succes = false;
	return sendBufIter;
}

RUDP::Fragment *RUDP::NewFragment(const char *src, unsigned int size)
{
	char *buf = new char[size];
	memcpy(buf, src, size);
	return new Fragment(buf, size);
}

unsigned int RUDP::Clock()
{
	return (std::chrono::time_point_cast<std::chrono::milliseconds>(std::chrono::system_clock::now()) - tp).count() & 0xffffffff;
}

bool RUDP::InRange(unsigned int v, unsigned int min, unsigned int max)
{
	return v <= max && v >= min;
}

// this is test function, is it stable
void Test()
{
	//std::queue<RUDP::UserData*>udpBuffer1;
	//std::queue<RUDP::UserData*>udpBuffer2;
	//std::queue<RUDP::UserData*> received_msg;

	//std::random_device rd;
	//std::mt19937 mt(rd());
	//std::uniform_int_distribution<int>dis(1, 100);

	//RUDP rudp1(100, 3, 512, 121);
	//RUDP rudp2(100, 3, 512, 121);

	//rudp1.SetOutput([&](const char* buf, unsigned short size)
	//	{
	//		if (dis(mt) >= 80)return;
	//		char* data = new char[size];
	//		std::memcpy(data, buf, size);
	//		udpBuffer2.push(new RUDP::UserData(data, size));
	//	});
	//rudp2.SetOutput([&](const char* buf, unsigned short size)
	//	{
	//		if (dis(mt) >= 70)return;
	//		char* data = new char[size];
	//		std::memcpy(data, buf, size);
	//		udpBuffer1.push(new RUDP::UserData(data, size));
	//	});

	//unsigned test = 0;
	//unsigned int counter = 0;
	//unsigned int locCounter = 0;
	//char send_msg[8];
	//bool run = false;
	//while (true)
	//{
	//	rudp1.Update();
	//	rudp2.Update();

	//	if (_kbhit() != 0) {
	//		run = !run;
	//		_getch();
	//	}
	//
	//	if (run)
	//	{
	//		sprintf_s(send_msg, sizeof(send_msg), "%d", counter++);
	//		rudp1.Send(send_msg, sizeof(send_msg));
	//	}

	//	//	receive message
	//	while (!udpBuffer2.empty())
	//	{
	//		rudp2.Input(udpBuffer2.front()->buf, udpBuffer2.front()->size);
	//		delete[] udpBuffer2.front()->buf;
	//		delete udpBuffer2.front();
	//		udpBuffer2.pop();
	//	}

	//	//	receive ack
	//	while (!udpBuffer1.empty())
	//	{
	//		rudp1.Input(udpBuffer1.front()->buf, udpBuffer1.front()->size);
	//		delete[] udpBuffer1.front()->buf;
	//		delete udpBuffer1.front();
	//		udpBuffer1.pop();
	//	}

	//	if (rudp2.Recv(received_msg))
	//	{
	//		RUDP::UserData* userd;
	//		while (!received_msg.empty())
	//		{
	//			userd = received_msg.front();
	//			received_msg.pop();

	//			unsigned int recvSeq = atoi(userd->buf);
	//			if (locCounter != recvSeq)
	//				printf("ERROR HIT\n");
	//			printf("%u\n", recvSeq);
	//			RUDP::releaseFragment((RUDP::Fragment*)userd);
	//			locCounter++;
	//		}
	//	}
	//}
}
